from ctypes import *
import numpy as np

from LTESync.sequence import PSS

class ScanPSS:
    def __init__(self,r=12):
        self.pss = [PSS(t) for t in range(3)]
        self.fd = np.zeros((3,9600),dtype=complex)
        for t in range(3):
            td = np.zeros(9600,dtype=complex)
            td[:128] = np.fft.ifft(self.pss[t].fd)
            self.fd[t,:] = np.conj(np.fft.fft(td))
        self.rg = np.arange(-r,r)*25
        self.rlen = len(self.rg)
    
    def fd_pss(self,t,s):
        fd = np.zeros(9600,dtype=complex)
        if s!=0:
            fd[s:]  = self.fd[t,:-s]
            fd[:s] = self.fd[t,-s:]
        else:
            fd[:] = self.fd[t,:]
        return fd

    def find_all_peak(self,buf,init_f):
        pos = np.zeros((3,self.rlen,9600))
        n = buf.shape[0]//9600-1
        for df in range(self.rlen):
            f = self.rg[df]*200
            drc = 1+f/init_f
            for t in range(3):
                for fn in range(n):
                    ppos = int((fn*9600)/drc)
                    dat_fft = np.fft.fft(buf[ppos:ppos+9600]) 
                    pss_fft = self.fd_pss(t,self.rg[df])
                    pos[t,df,:] += np.abs(np.fft.ifft(dat_fft*pss_fft))
        amax = pos.argmax()
        t = amax//(self.rlen*9600)
        df = amax%(self.rlen*9600)
        pos = df%9600
        df = df//9600
        pos += 9600-960
        pos %= 9600
        fr = np.dot(buf[128:],np.conj(buf[:-128]))
        f = np.angle(fr)/2/np.pi*15e3
        pf = self.rg[df]*200
        # print(pf,f)
        ddf = round((pf-f)/15e3)*15e3+f
        return t,ddf,pos

    def find_all_pss(self,buf):
        pos = np.zeros((3,9600))
        n = buf.shape[0]//9600-1
        for t in range(3):
            for fn in range(n):
                ppos = int(fn*9600)
                dat_fft = np.fft.fft(buf[ppos:ppos+9600]) 
                pss_fft = self.fd_pss(t,0)
                pos[t,:] += np.abs(np.fft.ifft(dat_fft*pss_fft))
        amax = pos.argmax()
        t = amax//9600
        pos = amax%9600
        pos += 9600-960
        pos %= 9600
        print(f"find pss: cid={t} pos={pos}")
        return t,pos